#!/usr/bin/env python
# encoding: utf-8

from tqdm import tqdm
import numpy as np
from .utils import compute_eer

def length_norm(vector):
    assert len(vector.shape) == 1
    dim = len(vector)
    norm = np.linalg.norm(vector)
    return np.sqrt(dim) * vector / norm


def cosine_score(trials, index_mapping, eval_vectors, apply_length_norm=False, score_save_path=None):
    labels = []
    scores = []
    if apply_length_norm:
        print("apply length norm")

    for item in trials:
        enroll_vector = eval_vectors[index_mapping[item[1]]]
        test_vector = eval_vectors[index_mapping[item[2]]]
        if apply_length_norm:
            enroll_vector = length_norm(enroll_vector)
            test_vector = length_norm(test_vector)
        score = enroll_vector.dot(test_vector.T)
        denom = np.linalg.norm(enroll_vector) * np.linalg.norm(test_vector)
        score = score/denom

        labels.append(int(item[0]))
        scores.append(score)
    # compute_eer sklearn implementation
    eer, threshold = compute_eer(labels, scores)
    if score_save_path is not None:
        with open(score_save_path, "w") as f:
            for i in range(len(labels)):
                f.write("{} {}\n".format(labels[i], scores[i]))
        print("score results save to: {}".format(score_save_path))
    return eer, threshold


def PLDA_score(trials, index_mapping, eval_vectors, plda_analyzer):
    labels = []
    scores = []
    for item in trials:
        enroll_vector = eval_vectors[index_mapping[item[1]]]
        test_vector = eval_vectors[index_mapping[item[2]]]
        score = plda_analyzer.NLScore(enroll_vector, test_vector)
        labels.append(int(item[0]))
        scores.append(score)
    # compute_eer sklearn implementation
    eer, threshold = compute_eer(labels, scores)
    return eer, th

